﻿#pragma once

#include <list>
#include "Containers/Vector.h"
#include "Containers/Map.h"
#include <vulkan/vulkan.hpp>

//不同的用途决定了底层内存池的大小
//允许精细控制内存池
enum class EVulkanMemoryUsage
{
    UniformBuffer,
    VertexBuffer,
    IndexBuffer,
    Texture,
    //始终拥有内存映射的缓冲区
    MappedTexture,
};


//一小块内存，是是内存池分配的最小单位
struct SVulkanMemoryBlock
{
public:
    vk::DeviceMemory Memory=nullptr;
    //在Memory中的偏移
    int64_t Offset=0;
    //实际需要的内存大小
    int64_t Size=0;
    //容量， >=实际需要的内存大小
    int64_t Capacity=0;

    //这一块对应的映射内存，仅部分内存块有映射的内存
    void* MappedMemory=nullptr;

    EVulkanMemoryUsage Usage;
    uint32_t MemoryTypeIndex=0;

    //所属的大块内存,仅用于释放内存
    //如果是-1，表示这块内存不是由内存池分配的
    int64_t ChunkIndex=-1;
    //查找大块内存的索引
    int64_t BlockSizeIndex=0;

    bool IsValid() const
    {
        return ChunkIndex!=-1;
    }

    void Reset()
    {
        Memory=nullptr;
        Offset=0;
        Size=0;
        Capacity=0;
        Usage=EVulkanMemoryUsage::UniformBuffer;
        MemoryTypeIndex=0;
        ChunkIndex=-1;
    }
};

//一大块内存，用于存储一些小块内存
//同一个大块中的小块内存大小是固定的
class CVulkanMemoryChunk
{
public:

    CVulkanMemoryChunk(
        uint32_t InMemoryTypeIndex
        ,vk::Device InDevice
        ,int64_t InTotalSize
        ,int64_t InBlockSize
        ,bool bAutoMap
        );

    ~CVulkanMemoryChunk();

    bool IsEmpty() const
    {
        return AllocatedBlockCount==0;
    }

    bool IsFull() const
    {
        return UnusedBlockOffsets.empty();
    }

    int64_t AllocateBlock()
    {
        assert(!IsFull()&&"内存块已满");
        
        int64_t Offset=UnusedBlockOffsets.front();
        UnusedBlockOffsets.pop_front();

        AllocatedBlockCount++;

        return Offset;
    }

    void DeallocateBlock(const SVulkanMemoryBlock& Block)
    {
        UnusedBlockOffsets.push_back(Block.Offset);

        AllocatedBlockCount--;
    }

    uint32_t MemoryTypeIndex;
    vk::DeviceMemory Memory;

    int64_t GetBlockSize() const
    {
        return BlockSize;
    }

    int64_t GetTotalSize() const
    {
        return TotalSize;
    }
protected:
    const int64_t TotalSize; //总大小(字节)
    const int64_t BlockSize; //每一小块内存大小(字节)

    std::list<int64_t> UnusedBlockOffsets;

    //已分配的小块内存的个数
    int64_t AllocatedBlockCount=0; 

public:
    bool HasMappedMemory() const
    {
        return MappedMemory!=nullptr;
    }

    uint8_t* GetMappedMemory() const
    {
        return MappedMemory;
    }

protected:
    //映射的内存
    uint8_t* MappedMemory=nullptr;
};

struct SMemoryPoolSizes
{
    SMemoryPoolSizes(int64_t InChunkSize,int64_t InMinBlockSize,int64_t InMaxBlockSize)
        :ChunkSize(InChunkSize)
        ,MinBlockSize(InMinBlockSize)
        ,MaxBlockSize(InMaxBlockSize)
    {
        assert(MinBlockSize<=MaxBlockSize&&"MinBlockSize必须小于等于MaxBlockSize");
        assert(ChunkSize>=MaxBlockSize&&"ChunkSize必须大于等于MaxBlockSize");
    }

    //每一大块内存的大小(字节)
    int64_t ChunkSize;
    //每一小块内存大小的下限(字节)
    int64_t MinBlockSize;
    //每一小块内存大小的上限(字节)
    int64_t MaxBlockSize;
};

//由于Vulkan的内存对象数量有限
//即vk::DeviceMemory的个数不大于VkPhysicalDeviceLimits::maxMemoryAllocationCount
//，所以需要一个内存池来管理内存对象
class CMemoryPool
{

public:

    CMemoryPool(
        EVulkanMemoryUsage InUsage
        , const SMemoryPoolSizes & InSizes
        ,uint32_t InMemoryTypeIndex);
    ~CMemoryPool(); 

    //分配一块内存
    SVulkanMemoryBlock AllocateBlock(const vk::MemoryRequirements& Requirement);

    //释放一块内存
    void DeallocateBlock(const SVulkanMemoryBlock& Block);
protected:

    //专门用于存储UniformBuffer的内存块,用于向着色器传递参数
    //它的SharingMode = vk::SharingMode::eExclusive
    //Chunks[BlockSize][ChunkIndex]
    TVector<TVector<CVulkanMemoryChunk*>> Chunks;

    const uint32_t MemoryTypeIndex;

    const EVulkanMemoryUsage Usage;

    SMemoryPoolSizes Sizes;
};



class CMemoryPools
{
public:
    CMemoryPools()=default;

    void Init();

    SVulkanMemoryBlock AllocateBlock(
        EVulkanMemoryUsage InUsage
        ,uint32_t InMemoryTypeIndex
        ,const vk::MemoryRequirements& Requirement
        );

    void DeallocateBlock(const SVulkanMemoryBlock& Block);
protected:

    struct SMemoryPoolKey
    {
        uint32_t MemoryTypeIndex;
        EVulkanMemoryUsage Usage;

        uint64_t GetHashCode() const
        {
            return MemoryTypeIndex+uint64_t(Usage)<<32;
        }

        bool operator==(const SMemoryPoolKey& InOther) const
        {
            return MemoryTypeIndex==InOther.MemoryTypeIndex&&Usage==InOther.Usage;
        }
    };

    TMap<SMemoryPoolKey ,std::unique_ptr<CMemoryPool>> MemoryPools;

    TMap<EVulkanMemoryUsage,SMemoryPoolSizes> PoolSizes;
    
};




